package com.gitee.empty_null.spark;

import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import okio.ByteString;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.http.codec.ServerSentEvent;
import reactor.core.publisher.FluxSink;

/**
 * @author xuhainan
 * @date 2023/11/17 17:07
 * @region hefei
 */
@Slf4j
public class SparkChatListener extends WebSocketListener {

    private final FluxSink<ServerSentEvent<String>> sink;
    private final String event;

    public SparkChatListener(FluxSink<ServerSentEvent<String>> sink, String event) {
        this.sink = sink;
        this.event = event;
    }

    @Override
    public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
        log.info("onClosed socket {} code {} reason {}", webSocket, code, reason);
        sink.complete();
    }

    @Override
    public void onClosing(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
        log.info("onClosing socket {} code {} reason {}", webSocket, code, reason);
        sink.complete();
    }

    @Override
    public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, @Nullable Response response) {
        log.info("onFailure socket {} t {} response {}", webSocket, t, response);
        sink.error(t);
    }

    @Override
    public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
        log.info("onMessage socket {} text {} ", webSocket, text);
        String content = SparkChatResponse.content(text);
        ServerSentEvent.Builder<String> builder = ServerSentEvent.builder();
        if (StrUtil.isNotBlank(event)) {
            builder.event(event);
        }
        ServerSentEvent<String> sentEvent = builder.data(content).build();
        sink.next(sentEvent);
    }

    @Override
    public void onMessage(@NotNull WebSocket webSocket, @NotNull ByteString bytes) {
        log.info("onMessage socket {} bytes {} ", webSocket, bytes);
        ServerSentEvent.Builder<String> builder = ServerSentEvent.builder();
        if (StrUtil.isNotBlank(event)) {
            builder.event(event);
        }
        ServerSentEvent<String> sentEvent = builder.data(bytes.toString()).build();
        sink.next(sentEvent);
    }

    @Override
    public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
        log.info("onOpen socket {} response {} ", webSocket, response);
    }
}
